import weka.core.converters.ConverterUtils.DataSource;
import weka.core.Instances;
import weka.classifiers.trees.RandomForest;
import weka.classifiers.Evaluation;

import java.io.File;
import java.io.FileOutputStream;
import java.io.FileInputStream;
import java.io.ObjectOutputStream;
import java.io.ObjectInputStream;
import java.io.PrintWriter;
import java.util.Arrays;

public class WekaRandomForest
{
    /*
     * Build Classification model
     * args[0] = train dataset input file
     * args[1] = object output file
     * args[2] = Number of tree in forest
     * args[3] = Number of features to consider
     * args[4] = Random number seed
     * args[5] = Maximum depth
     */
    static void build(String[] args) throws Exception {
        // Train the classifier
        DataSource source = new DataSource(args[0]);
        Instances data = source.getDataSet();
        data.setClassIndex(data.numAttributes() - 1);

        RandomForest tree = new RandomForest();
        tree.setNumTrees(Integer.parseInt(args[2]));
        tree.setNumFeatures(Integer.parseInt(args[3]));
        tree.setSeed(Integer.parseInt(args[4]));
        tree.setMaxDepth(Integer.parseInt(args[5]));
        tree.buildClassifier(data);   // build classifier

        // Serialize the classifier
        ObjectOutputStream oos = new ObjectOutputStream(
            new FileOutputStream(new File(args[1]))
        );
        oos.writeObject(tree);
        oos.close();
    }

    /*
     * Test Classification model
     * args[0] = test data file
     * args[1] = RandomForest object file
     * args[2] = Result output file
     */
    static void test(String[] args) throws Exception {
        // Deserialize the classifier
        RandomForest tree = new RandomForest();         // new instance of tree

        ObjectInputStream ois = new ObjectInputStream(
            new FileInputStream(new File(args[1]))
        );
        tree = (RandomForest) ois.readObject();
        ois.close();

        // Test the classifier
        DataSource source = new DataSource(args[0]);
        Instances data = source.getDataSet();
        data.setClassIndex(data.numAttributes() - 1);

        Evaluation eval = new Evaluation(data);
        eval.useNoPriors();
        eval.evaluateModel(tree, data);

        // Output the result
        PrintWriter pw = new PrintWriter(args[2]);
        pw.println(eval.toSummaryString("\nResults\n======\n", false));
        pw.close();
    }

    /*
     * Classify some records
     * args[0] = Dataset input file
     * args[1] = RandomForest object file
     * args[2] = Result output file
     */
    static void classify(String[] args) throws Exception {
        // Deserialize the classifier
        RandomForest tree = new RandomForest();         // new instance of tree

        ObjectInputStream ois = new ObjectInputStream(
            new FileInputStream(new File(args[1]))
        );
        tree = (RandomForest) ois.readObject();
        ois.close();

        // Get output
        PrintWriter pw = new PrintWriter(args[2]);

        DataSource source = new DataSource(args[0]);
        Instances data = source.getDataSet();
        data.setClassIndex(data.numAttributes() - 1);

        for(int i = 0; i < data.numInstances(); ++i) {
            data.instance(i).setClassValue(-1.0);
            double guess = tree.classifyInstance(data.instance(i));

            pw.println(i + " " + data.classAttribute().value((int) guess));
        }

        pw.close();
    }

    public static void main(String[] args) throws Exception {
        String task = args[0];
        args = Arrays.copyOfRange(args, 1, args.length);

        System.out.println("Call Weka Random Forest " + task);

        if (task.equals("build")) {
            build(args);
        }
        else if (task.equals("test")) {
            test(args);
        }
        else if (task.equals("classify")) {
            classify(args);
        }
    }
}

// Compile: javac -cp /home/rr/DataMining/WEKA/weka-3-6-10/weka.jar WekaRandomForest.java
// Run: java -Xmx2108m -cp /home/rr/DataMining/WEKA/weka-3-6-10/weka.jar:. WekaRandomForest
